{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1 align=\"center\">Introduction to SimpleITKv4 Registration - Continued</h1>\n",
    "\n",
    "\n",
    "## ITK v4 Registration Components \n",
    "<img src=\"ITKv4RegistrationComponentsDiagram.svg\" style=\"width:700px\"/><br><br>\n",
    "\n",
    "\n",
    "Before starting with this notebook, please go over the first introductory notebook found [here](60_Registration_Introduction.ipynb). \n",
    "\n",
    "In this notebook we will visually assess registration by viewing the overlap between images using external viewers.\n",
    "The two viewers we recommend for this task are [ITK-SNAP](http://www.itksnap.org) and [3D Slicer](http://www.slicer.org/). ITK-SNAP supports concurrent linked viewing between multiple instances of the program. 3D Slicer supports concurrent viewing of multiple volumes via alpha blending."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "\n",
    "# If the environment variable SIMPLE_ITK_MEMORY_CONSTRAINED_ENVIRONMENT is set, this will override the ReadImage\n",
    "# function so that it also resamples the image to a smaller size (testing environment is memory constrained).\n",
    "%run setup_for_testing\n",
    "\n",
    "# Utility method that either downloads data from the network or\n",
    "# if already downloaded returns the file name for reading from disk (cached data).\n",
    "%run update_path_to_download_script\n",
    "from downloaddata import fetch_data as fdata\n",
    "\n",
    "# Always write output to a separate directory, we don't want to pollute the source directory. \n",
    "import os\n",
    "OUTPUT_DIR = 'Output'\n",
    "\n",
    "# GUI components (sliders, dropdown...).\n",
    "from ipywidgets import interact, fixed\n",
    "     \n",
    "# Enable display of HTML.\n",
    "from IPython.display import display, HTML \n",
    "    \n",
    "# Plots will be inlined.\n",
    "%matplotlib inline\n",
    "\n",
    "# Callbacks for plotting registration progress.\n",
    "import registration_callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility functions\n",
    "A number of utility functions, saving a transform and corresponding resampled image, callback for selecting a \n",
    "DICOM series from several series found in the same directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_transform_and_image(transform, fixed_image, moving_image, outputfile_prefix):\n",
    "    \"\"\"\n",
    "    Write the given transformation to file, resample the moving_image onto the fixed_images grid and save the\n",
    "    result to file.\n",
    "    \n",
    "    Args:\n",
    "        transform (SimpleITK Transform): transform that maps points from the fixed image coordinate system to the moving.\n",
    "        fixed_image (SimpleITK Image): resample onto the spatial grid defined by this image.\n",
    "        moving_image (SimpleITK Image): resample this image.\n",
    "        outputfile_prefix (string): transform is written to outputfile_prefix.tfm and resampled image is written to \n",
    "                                    outputfile_prefix.mha.\n",
    "    \"\"\"                             \n",
    "    resample = sitk.ResampleImageFilter()\n",
    "    resample.SetReferenceImage(fixed_image)\n",
    "    \n",
    "    # SimpleITK supports several interpolation options, we go with the simplest that gives reasonable results.     \n",
    "    resample.SetInterpolator(sitk.sitkLinear)  \n",
    "    resample.SetTransform(transform)\n",
    "    sitk.WriteImage(resample.Execute(moving_image), outputfile_prefix+'.mha')\n",
    "    sitk.WriteTransform(transform, outputfile_prefix+'.tfm')\n",
    "    \n",
    "def DICOM_series_dropdown_callback(fixed_image, moving_image, series_dictionary):\n",
    "    \"\"\"\n",
    "    Callback from dropbox which selects the two series which will be used for registration.\n",
    "    The callback prints out some information about each of the series from the meta-data dictionary.\n",
    "    For a list of all meta-dictionary tags and their human readable names see DICOM standard part 6, \n",
    "    Data Dictionary (http://medical.nema.org/medical/dicom/current/output/pdf/part06.pdf)\n",
    "    \"\"\"\n",
    "    # The callback will update these global variables with the user selection.\n",
    "    global selected_series_fixed\n",
    "    global selected_series_moving\n",
    "    \n",
    "    img_fixed = sitk.ReadImage(series_dictionary[fixed_image][0])\n",
    "    img_moving = sitk.ReadImage(series_dictionary[moving_image][0])\n",
    "    \n",
    "    # There are many interesting tags in the DICOM data dictionary, display a selected few.\n",
    "    tags_to_print = {'0010|0010': 'Patient name: ', \n",
    "                     '0008|0060' : 'Modality: ',\n",
    "                     '0008|0021' : 'Series date: ',\n",
    "                     '0008|0031' : 'Series time:',\n",
    "                     '0008|0070' : 'Manufacturer: '}\n",
    "    html_table = []\n",
    "    html_table.append('<table><tr><td><b>Tag</b></td><td><b>Fixed Image</b></td><td><b>Moving Image</b></td></tr>')\n",
    "    for tag in tags_to_print:\n",
    "        fixed_tag = ''\n",
    "        moving_tag = ''\n",
    "        try:            \n",
    "            fixed_tag = img_fixed.GetMetaData(tag)\n",
    "        except: # ignore if the tag isn't in the dictionary\n",
    "            pass\n",
    "        try:            \n",
    "            moving_tag = img_moving.GetMetaData(tag)\n",
    "        except: # ignore if the tag isn't in the dictionary\n",
    "            pass           \n",
    "        html_table.append('<tr><td>' + tags_to_print[tag] + \n",
    "                          '</td><td>' + fixed_tag + \n",
    "                          '</td><td>' + moving_tag + '</td></tr>')\n",
    "    html_table.append('</table>')\n",
    "    display(HTML(''.join(html_table)))\n",
    "    selected_series_fixed = fixed_image\n",
    "    selected_series_moving = moving_image    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading Data\n",
    "\n",
    "In this notebook we will work with CT and MR scans of the CIRS 057A multi-modality abdominal phantom. The scans are multi-slice DICOM images. The data is stored in a zip archive which is automatically retrieved and extracted when we request a file which is part of the archive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_directory = os.path.dirname(fdata(\"CIRS057A_MR_CT_DICOM/readme.txt\"))\n",
    "\n",
    "# 'selected_series_moving/fixed' will be updated by the interact function.\n",
    "selected_series_fixed = ''\n",
    "selected_series_moving = ''\n",
    "\n",
    "# Directory contains multiple DICOM studies/series, store the file names\n",
    "# in dictionary with the key being the series ID.\n",
    "reader = sitk.ImageSeriesReader()\n",
    "series_file_names = {}\n",
    "series_IDs = list(reader.GetGDCMSeriesIDs(data_directory)) #list of all series\n",
    "            \n",
    "if series_IDs: #check that we have at least one series\n",
    "    for series in series_IDs:\n",
    "        series_file_names[series] = reader.GetGDCMSeriesFileNames(data_directory, series)    \n",
    "    interact(DICOM_series_dropdown_callback, fixed_image=series_IDs, moving_image =series_IDs, series_dictionary=fixed(series_file_names)); \n",
    "else:\n",
    "    print('This is surprising, data directory does not contain any DICOM series.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Actually read the data based on the user's selection.\n",
    "fixed_image = sitk.ReadImage(series_file_names[selected_series_fixed])\n",
    "moving_image = sitk.ReadImage(series_file_names[selected_series_moving])\n",
    "\n",
    "# Save images to file and view overlap using external viewer.\n",
    "sitk.WriteImage(fixed_image, os.path.join(OUTPUT_DIR, \"fixedImage.mha\"))\n",
    "sitk.WriteImage(moving_image, os.path.join(OUTPUT_DIR, \"preAlignment.mha\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial Alignment\n",
    "\n",
    "A reasonable guesstimate for the initial translational alignment can be obtained by using\n",
    "the CenteredTransformInitializer (functional interface to the CenteredTransformInitializerFilter). \n",
    "\n",
    "The resulting transformation is centered with respect to the fixed image and the\n",
    "translation aligns the centers of the two images. There are two options for\n",
    "defining the centers of the images, either the physical centers\n",
    "of the two data sets (GEOMETRY), or the centers defined by the intensity \n",
    "moments (MOMENTS).\n",
    "\n",
    "Two things to note about this filter, it requires the fixed and moving image \n",
    "have the same type even though it is not algorithmically required, and its\n",
    "return type is the generic SimpleITK.Transform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_transform = sitk.CenteredTransformInitializer(sitk.Cast(fixed_image,moving_image.GetPixelID()), \n",
    "                                                      moving_image, \n",
    "                                                      sitk.Euler3DTransform(), \n",
    "                                                      sitk.CenteredTransformInitializerFilter.GEOMETRY)\n",
    "\n",
    "# Save moving image after initial transform and view overlap using external viewer.\n",
    "save_transform_and_image(initial_transform, fixed_image, moving_image, os.path.join(OUTPUT_DIR, \"initialAlignment\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the transformation, what type is it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(initial_transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final registration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Version 1\n",
    "<ul>\n",
    "<li> Single scale (not using image pyramid).</li>\n",
    "<li> Initial transformation is not modified in place.</li>\n",
    "</ul>\n",
    "\n",
    "<ol>\n",
    "<li>\n",
    "Illustrate the need for scaling the step size differently for each parameter:\n",
    "<ul>\n",
    "<li> SetOptimizerScalesFromIndexShift - estimated from maximum shift of voxel indexes (only use if data is isotropic).</li>\n",
    "<li> SetOptimizerScalesFromPhysicalShift - estimated from maximum shift of physical locations of voxels.</li>\n",
    "<li> SetOptimizerScalesFromJacobian - estimated from the averaged squared norm of the Jacobian w.r.t. parameters.</li>\n",
    "</ul>\n",
    "</li>\n",
    "<li>\n",
    "Look at the optimizer's stopping condition to ensure we have not terminated prematurely.\n",
    "</li>\n",
    "</ol>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "registration_method = sitk.ImageRegistrationMethod()\n",
    "\n",
    "registration_method.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)\n",
    "registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)\n",
    "registration_method.SetMetricSamplingPercentage(0.01)\n",
    "\n",
    "registration_method.SetInterpolator(sitk.sitkLinear)\n",
    "\n",
    "registration_method.SetOptimizerAsGradientDescent(learningRate=1.0, numberOfIterations=100)\n",
    "# Scale the step size differently for each parameter, this is critical!!!\n",
    "registration_method.SetOptimizerScalesFromPhysicalShift() \n",
    "\n",
    "registration_method.SetInitialTransform(initial_transform, inPlace=False)\n",
    "\n",
    "registration_method.AddCommand(sitk.sitkStartEvent, registration_callbacks.metric_start_plot)\n",
    "registration_method.AddCommand(sitk.sitkEndEvent, registration_callbacks.metric_end_plot)\n",
    "registration_method.AddCommand(sitk.sitkIterationEvent, \n",
    "                               lambda: registration_callbacks.metric_plot_values(registration_method))\n",
    "\n",
    "final_transform_v1 = registration_method.Execute(sitk.Cast(fixed_image, sitk.sitkFloat32), \n",
    "                                                 sitk.Cast(moving_image, sitk.sitkFloat32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Optimizer\\'s stopping condition, {0}'.format(registration_method.GetOptimizerStopConditionDescription()))\n",
    "print('Final metric value: {0}'.format(registration_method.GetMetricValue()))\n",
    "\n",
    "# Save moving image after registration and view overlap using external viewer.\n",
    "save_transform_and_image(final_transform_v1, fixed_image, moving_image, os.path.join(OUTPUT_DIR, \"finalAlignment-v1\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the final transformation, what type is it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(final_transform_v1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Version 1.1\n",
    "\n",
    "The previous example illustrated the use of the ITK v4 registration framework in an ITK v3 manner. We only referred to a single transformation which was what we optimized.\n",
    "\n",
    "In ITK v4 the registration method accepts three transformations (if you look at the diagram above you will only see two transformations, Moving transform represents $T_{opt} \\circ T_m$):\n",
    "<ul>\n",
    "<li>\n",
    "SetInitialTransform, $T_{opt}$ - composed with the moving initial transform, maps points from the virtual image domain to the moving image domain, modified during optimization. \n",
    "</li>\n",
    "<li>\n",
    "SetFixedInitialTransform $T_f$- maps points from the virtual image domain to the fixed image domain, never modified.\n",
    "</li>\n",
    "<li>\n",
    "SetMovingInitialTransform $T_m$- maps points from the virtual image domain to the moving image domain, never modified.\n",
    "</li>\n",
    "</ul>\n",
    "\n",
    "The transformation that maps points from the fixed to moving image domains is thus: $^M\\mathbf{p}  = T_{opt}(T_m(T_f^{-1}(^F\\mathbf{p})))$\n",
    "\n",
    "We now modify the previous example to use $T_{opt}$ and $T_m$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "registration_method = sitk.ImageRegistrationMethod()\n",
    "registration_method.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)\n",
    "registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)\n",
    "registration_method.SetMetricSamplingPercentage(0.01)\n",
    "registration_method.SetInterpolator(sitk.sitkLinear)\n",
    "registration_method.SetOptimizerAsGradientDescent(learningRate=1.0, numberOfIterations=100)\n",
    "registration_method.SetOptimizerScalesFromPhysicalShift()\n",
    "\n",
    "# Set the initial moving and optimized transforms.\n",
    "optimized_transform = sitk.Euler3DTransform()    \n",
    "registration_method.SetMovingInitialTransform(initial_transform)\n",
    "registration_method.SetInitialTransform(optimized_transform)\n",
    "\n",
    "registration_method.AddCommand(sitk.sitkStartEvent, registration_callbacks.metric_start_plot)\n",
    "registration_method.AddCommand(sitk.sitkEndEvent, registration_callbacks.metric_end_plot)\n",
    "registration_method.AddCommand(sitk.sitkIterationEvent, \n",
    "                               lambda: registration_callbacks.metric_plot_values(registration_method))\n",
    "\n",
    "registration_method.Execute(sitk.Cast(fixed_image, sitk.sitkFloat32), \n",
    "                            sitk.Cast(moving_image, sitk.sitkFloat32))\n",
    "\n",
    "# Need to compose the transformations after registration.\n",
    "final_transform_v11 = sitk.CompositeTransform(optimized_transform)\n",
    "final_transform_v11.AddTransform(initial_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Optimizer\\'s stopping condition, {0}'.format(registration_method.GetOptimizerStopConditionDescription()))\n",
    "print('Final metric value: {0}'.format(registration_method.GetMetricValue()))\n",
    "\n",
    "# Save moving image after registration and view overlap using external viewer.\n",
    "save_transform_and_image(final_transform_v11, fixed_image, moving_image, os.path.join(OUTPUT_DIR, \"finalAlignment-v1.1\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the final transformation, what type is it? Why is it different from the previous example?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(final_transform_v11)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Version 2\n",
    "\n",
    "<ul>\n",
    "<li> Multi scale - specify both scale, and how much to smooth with respect to original image.</li>\n",
    "<li> Initial transformation modified in place, so in the end we have the same type of transformation in hand.</li>\n",
    "</ul>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "registration_method = sitk.ImageRegistrationMethod()\n",
    "\n",
    "registration_method.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)\n",
    "registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)\n",
    "registration_method.SetMetricSamplingPercentage(0.01)\n",
    "\n",
    "registration_method.SetInterpolator(sitk.sitkLinear)\n",
    "   \n",
    "registration_method.SetOptimizerAsGradientDescent(learningRate=1.0, numberOfIterations=100) #, estimateLearningRate=registration_method.EachIteration)\n",
    "registration_method.SetOptimizerScalesFromPhysicalShift() \n",
    "\n",
    "final_transform = sitk.Euler3DTransform(initial_transform)\n",
    "registration_method.SetInitialTransform(final_transform)\n",
    "registration_method.SetShrinkFactorsPerLevel(shrinkFactors = [4,2,1])\n",
    "registration_method.SetSmoothingSigmasPerLevel(smoothingSigmas = [2,1,0])\n",
    "registration_method.SmoothingSigmasAreSpecifiedInPhysicalUnitsOn()\n",
    "\n",
    "registration_method.AddCommand(sitk.sitkStartEvent, registration_callbacks.metric_start_plot)\n",
    "registration_method.AddCommand(sitk.sitkEndEvent, registration_callbacks.metric_end_plot)\n",
    "registration_method.AddCommand(sitk.sitkMultiResolutionIterationEvent, \n",
    "                               registration_callbacks.metric_update_multires_iterations) \n",
    "registration_method.AddCommand(sitk.sitkIterationEvent, \n",
    "                               lambda: registration_callbacks.metric_plot_values(registration_method))\n",
    "\n",
    "registration_method.Execute(sitk.Cast(fixed_image, sitk.sitkFloat32), \n",
    "                            sitk.Cast(moving_image, sitk.sitkFloat32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Optimizer\\'s stopping condition, {0}'.format(registration_method.GetOptimizerStopConditionDescription()))\n",
    "print('Final metric value: {0}'.format(registration_method.GetMetricValue()))\n",
    "\n",
    "# Save moving image after registration and view overlap using external viewer.\n",
    "save_transform_and_image(final_transform, fixed_image, moving_image, os.path.join(OUTPUT_DIR, 'finalAlignment-v2'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the final transformation, what type is it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(final_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
